import os
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer

from src.database.dao.video_vector_mapping_dao import VideoVectorMappingDAO
from src.database.models import VideoVectorMapping

class FAISSProcessor:
    def __init__(self, index_path='videos.index', model_name='shibing624/text2vec-base-multilingual'):
        self.embedder = SentenceTransformer(model_name)
        self.dimension = len(self.embedder.encode("测试语句"))
        self.index = faiss.IndexIDMap2(faiss.IndexFlatL2(self.dimension))
        self.index_path = index_path
        if os.path.exists(self.index_path):
            print(f"Loading index from {self.index_path}")
            self.index = faiss.read_index(self.index_path)
        else:
            print(f"Creating new index at dimension {self.dimension}")
            self.index = faiss.IndexIDMap2(faiss.IndexFlatL2(self.dimension))
            self.save_index()
        
    def create_index(self, vectors):
        """Create a new FAISS index from vectors"""
        if not isinstance(vectors, np.ndarray):
            vectors = np.array(vectors).astype('float32')
        if len(vectors.shape) == 1:
            vectors = vectors.reshape(1, -1)
        self.index.add(vectors)
        return self
        
    def search(self, text, k=3):
        """Search for the top k nearest neighbors using text input and return video_ids"""
        # Convert text to vector using the embedder
        text = f"一个关于{text}的视频"
        query_vector = self.embedder.encode(text)
        query_vector = np.array(query_vector).astype('float32')
        if len(query_vector.shape) == 1:
            query_vector = query_vector.reshape(1, -1)
        
        # Get distances and FAISS indices
        distances, indices = self.index.search(query_vector, k)
        
        # Map FAISS indices to video_ids
        results = []
        dao = VideoVectorMappingDAO()
        for i in range(len(indices[0])):
            faiss_id = indices[0][i]
            distance = distances[0][i]
            print(f"Processing index {i}, faiss_id: {faiss_id}, distance: {distance}")
            video_id = dao.get_video_id_by_faiss_id(faiss_id)
            results.append((distance, video_id))
        
        return results
    
    def get_index_count(self):
        """Get the total number of vectors in the FAISS index"""
        return self.index.ntotal
    
    def save_index(self):
        faiss.write_index(self.index, self.index_path)
    def load_index(self):
        self.index = faiss.read_index(self.index_path)
        return self

    def add_video_description(self, video_id, description, 
                              username, category,
                              vector_type='text', frame_number=None):
        parts = description.split('#', 1)
        title = parts[0].strip()
        tags = []
        if len(parts) > 1:
            tags = [tag.strip() for tag in parts[1].split('#') if tag.strip()]
        encode_text = f"""这是一个视频，视频标题: {title}
        视频标签: {tags}
        视频作者: {username}
        视频分类: {category}"""
        vector = self.embedder.encode(encode_text)
        vector_np = np.array([vector]).astype('float32')

        print("Vector shape:", vector_np.shape)
        print("Index dimension:", self.index.d)
        print("Index type:", type(self.index))
        print("Index is_trained:", getattr(self.index, 'is_trained', True))  # Flat index 没这个属性
        print("Before add, ntotal:", self.index.ntotal)

        self.index.add_with_ids(vector_np, np.array([video_id]).astype('int64'))

        print("After add, ntotal:", self.index.ntotal)

        self.save_index()

        VideoVectorMappingDAO().add(
            video_id=video_id,
            faiss_id=video_id,
            vector_type=vector_type,
            frame_number=frame_number,
            vector_dimension=self.dimension,
        )
